import json
import PyPDF2
import importlib
from io import BytesIO

# 尝试导入python-docx模块
try:
    docx_module = importlib.import_module('docx')
    Document = getattr(docx_module, 'Document', None)
    docx_supported = True
except ImportError:
    Document = None
    docx_supported = False


def extract_code_cells(nb):
    return [line for c in nb.get('cells', [])
            if c.get('cell_type') == 'code'
            for line in c.get('source', [])
            if line.strip()]


def smart_extract(fs, filename):
    """
    智能提取文件内容
    :param fs: 文件流对象
    :param filename: 文件名
    :return: (内容列表, 文件类型, 语言)
    """
    try:
        # 获取文件扩展名
        ext = filename.rsplit('.', 1)[-1].lower() if '.' in filename else ''

        # 重置文件指针并读取内容
        fs.seek(0)
        content = fs.read()

        # 根据文件类型处理内容
        if ext == 'ipynb':
            try:
                notebook = json.loads(content.decode('utf-8'))
                return extract_code_cells(notebook), 'Notebook', 'python'
            except Exception as e:
                return f"处理Notebook文件错误: {str(e)}"

        if ext == 'docx' and docx_supported and Document:
            try:
                doc = Document(BytesIO(content))
                return [p.text for p in doc.paragraphs], 'Text', 'text'
            except Exception as e:
                return f"处理DOCX文件错误: {str(e)}"

        if ext == 'pdf':
            try:
                pdf = PyPDF2.PdfReader(BytesIO(content))
                return [line for page in pdf.pages
                        for line in (page.extract_text() or '').splitlines()], 'Text', 'text'
            except Exception as e:
                return f"处理PDF文件错误: {str(e)}"

        # 默认处理为文本
        try:
            return content.decode('utf-8', 'ignore').splitlines(), 'Text', 'text'
        except:
            return str(content).splitlines(), 'Text', 'text'

    except Exception as e:
        print(f"处理文件错误: {str(e)}")
        return [], 'Error', 'text'
